import requests
import json
import hashlib
import base64
from pathlib import Path
import subprocess
from uuid import uuid4
import numpy as np
import random
import os
import logging 
import base64
import re
import sys
#pip install JayDeBeApi
#import jaydebeapi
#import psycopg2
import psycopg2_pool
from langchain_core.tools import StructuredTool
from string import Template
from ApiEmbed import emObj

# pip install python-dotenv
from dotenv import load_dotenv
load_dotenv()

class ApiBase():
    #api_url: str = os.environ["APITOOLS_BASE_URL"]
    #api_url: str = "http://10.220.130.88:11410/"
    dir: str = os.path.dirname(os.path.abspath(__file__))    
    
    # 保证单实例
    _instance = None
    def __new__(cls, *args, **kwargs):
        if cls._instance is None:
            cls._instance = super(ApiBase, cls).__new__(cls, *args, **kwargs)
        return cls._instance
    
    def __init__(self):
        print("ApiBase init")
        logging.basicConfig(filename="/tmp/llmdb.log", level=logging.INFO)
        self.pool=None
        self.sys_dict={}
        self.lsExpress=[]
        self.funs_run={}
        self.agent_list={}
        self.tools_list={}
        #print("apiBase init")
    
    def close(self):
        self.pool.clear()
    
    def getLLMConnect(self):
        if self.pool:
            return self.pool
        ip=os.environ["LLMDB_IP"]
        port=os.environ["LLMDB_PORT"]
        #url=f'''dbname=./llm.data;DATABASE_TO_LOWER=TRUE;DEFAULT_NULL_ORDERING=HIGH;AUTO_SERVER=TRUE user=sa password='sa' host={ip} port={port} options='-c search_path=public' '''
        #self.db_connection = psycopg2.connect(url)
        self.pool = psycopg2_pool.ConnectionPool(minconn=1, maxconn=3,
                user="sa",
                password="sa",
                host=ip,
                port=port,
                options="-c search_path=public",
                database="./llm.data;DATABASE_TO_LOWER=TRUE;DEFAULT_NULL_ORDERING=HIGH;")
        #self.query_dict()
        return self.pool
    
    # def getLLMConnectJar(self):
    #     if self.db_connection is not None:
    #         return self.db_connection        
    #     #jdbc_url  =  os.environ["LLMDB_URL"]+';SCHEMA=public'
    #     jdbc_url  =  os.environ["LLMDB_URL"]
    #     driver_name = "org.h2.Driver"
    #     jar_path = os.environ["LLMDB_JAR_URL"]
    #     # 创建JDBC连接参数
    #     jdbc_params = {
    #         "driver": driver_name,
    #         "url": jdbc_url,
    #         "jar": jar_path,
    #         "SCHEMA":"public"
    #     }
    #     # 使用jaydebeapi建立连接并执行SQL查询
    #     db_connection = jaydebeapi.connect(jdbc_params['driver'],jdbc_params['url'],['sa','sa'],jdbc_params['jar'])
    #     self.db_connection=db_connection 
    #     #self.query_dict()
    #     return db_connection
    
    # 给出一个topn=9个的值，从topn里面 按权重生成数组,目前用于rag里面随机找出对话的向量，用于生成
    def get_random(self,topn, k):
        def weighted_random_choice(weights, k):
            # 计算总的权重
            total_weight = sum(weights)
            # 根据权重生成选择概率
            probabilities = [w / total_weight for w in weights]            
            # 使用 numpy 的 choice 函数进行加权随机选择
            chosen_indices = np.random.choice(len(weights), size=k, replace=False, p=probabilities)
            return chosen_indices
        if k > topn:
            k = topn
        # 定义从 0 到 9 的数字和其对应的权重（数字越小，权重越大）
        numbers = list(range(topn))
        weights = [topn - i for i in numbers]  # 权重: [10, 9, 8, 7, 6, 5, 4, 3, 2, 1]
        # 随机选择 6 个不同的数字
        chosen_numbers = [numbers[i] for i in weighted_random_choice(weights, k)]
        chosen_numbers.sort()  # 排序以便查看
        # print("随机选择的 6 个不同的数字是:", chosen_numbers)
        return chosen_numbers

    # 读取数据库的数据
    def query_dict(self):
        results=self.query_sql("select pkey,val from aigc_dict")
        for row in results:
            pkey,val = row
            self.sys_dict[pkey]=val
            
    # 从字典获取数据
    def dict(self,pkey, defVal=None):
        if pkey in self.sys_dict:
            return self.sys_dict[pkey]
        return defVal
    
    def execute(self,sqls,params=None,ret=False):
        """执行sql并返回ID"""     
        con=self.pool.getconn()  
        try:
            with con.cursor() as cursor:
                try:
                    if params:
                        for sql, param in zip(sqls, params):
                            cursor.execute(sql,param)                            
                    else:
                        for sql in sqls:
                            cursor.execute(sql)                            
                    # 提交事务
                    con.commit()
                    return
                except Exception as e:
                    if ret:
                        return str(e)
                #如果只是一条数据，就退出；如果是多条记录，那就每天单独执行，保证正确
                if len(sqls) == 1:
                    return None
                if params:
                    for sql, param in zip(sqls, params):
                        try:                    
                            cursor.execute(sql,param)   
                            # 提交事务
                            con.commit()
                        except Exception as e:
                            if ret:
                                return str(e)
                    return None
                for sql in sqls:
                    try:
                        cursor.execute(sql)
                        con.commit()
                    except Exception as e:
                        print(f"execute sql error:{str(e)}")
                        if ret:
                            return str(e)
                return None
        finally:
            self.pool.putconn(con)
    # 读取数据库的数据
    def query_sql(self,sql):
        con=self.pool.getconn()
        try:
            with con.cursor() as cur:
                cur.execute(sql)
                results = cur.fetchall()
                return results
        finally:
            self.pool.putconn(con)
    # 把游标对象转换为列表
    def convert_cursor(self,cursor):
        # 获取字段名称
        fields = [description[0] for description in cursor.description]
        # 获取所有行的数据
        rows = cursor.fetchall()
        data = []
        # 打印每一行的数据
        for row in rows:
            item={}
            for index, field in enumerate(fields):
                item[field] = row[index]
            data.append(item)
        return data
    
    def test_sql(self,sql,commit=None):
        self.getLLMConnect()        
        data=self.query_json(self,sql,commit)
        print(data)
        
    def query_json(self,sql,commit=None):
        con=self.pool.getconn()
        try:
            with con.cursor() as cur:
                cur.execute(sql)
                if commit:
                    con.commit()
                    return ""
                data=self.convert_cursor(cur)
            return data
        finally:
            self.pool.putconn(con)
    
    # 把游标对象转换为列表
    def convert_cursor(self,cursor):
        # 获取字段名称
        fields = [description[0] for description in cursor.description]
        # 获取所有行的数据
        rows = cursor.fetchall()
        data = []
        # 打印每一行的数据
        for row in rows:
            item={}
            for index, field in enumerate(fields):
                item[field] = row[index]
            data.append(item)
        return data    
    
    # 把env文件写入数据库
    def write_env(self,jsEnv):
        env =json.loads(jsEnv)
        url=env["LLMDB_URL"]
        os.environ["LLMDB_URL"]=env["LLMDB_URL"]
        os.environ["LLMDB_JAR_URL"]=env["LLMDB_JAR_URL"]
        self.getLLMConnect()
        con=self.pool.getconn()
        try:
            with con.cursor() as cur:
                sql = "SELECT DATABASE_PATH()"
                cur.execute(sql)
                results = cur.fetchall()
                for row in results:
                    dbpath, = row
                dir = os.path.dirname(dbpath)
                with open(dir+'/.env', 'w') as env_file:                
                    for key in env:
                        value = env[key]
                        env_file.write(f'{key}={value}\n')                    
        finally:
            self.pool.putconn(con)
        return "save ok" 
    
    def argv(self,funs_param,index,defVal=None):
        if len(funs_param) <= index:
            val=defVal
        else:
            val=apiBase.decode64(funs_param[index])
        if val == None:
            return val
        val=val.strip()
        #去掉字符串传过来的''字符
        if val.startswith("'") and val.endswith("'"):
            val=val[1:-1]
        if len(val) == 0:
            return defVal
        if val.find('$PROJECT_HOME') != -1:
            template = Template(val)
            val = template.substitute(PROJECT_HOME=os.environ["PROJECT_HOME"])
        return val        
    
    def argv_json(self,funs_param,index,defVal=None):
        try:
            if len(funs_param) <= index:
                data=defVal
            else:
                data=self.argv(funs_param,index,'')
            return json.loads(data)
        except Exception as e:
            return None

    def argv_int(self,funs_param,index,defVal=None):        
        try:
            if len(funs_param) <= index:
                data= defVal
            else:
                data=self.argv(funs_param,index,'')
            return int(data)
        except Exception as e:
            return defVal
        
    # 字符串转成函数
    def str_func(self,func):      
        try:
            # 提取函数名称
            md5=self.generate_md5(func)
            if md5 in self.tools_list:
                ret=self.tools_list[md5]
                return func
            pattern = r'def\s+(\w+)\('
            match = re.search(pattern, func) 
            func_name = match.group(1)                
            exec(func)
            ret=eval(f'{func_name}')
            self.tools_list[md5]=ret           
            return ret
        except Exception as e:
            return None
    
    # 动态创建langgraph工具
    def create_tool(self,name: str,desc: str,url: str,pip:str) -> dict:
        """Create schema for a placeholder tool."""
        def company_tool(param: str) -> str:
            if pip == 'hpl':
                ret=apiBase.call_api(url,param)
            if pip == 'rest':
                ret=apiBase.call_pip(url,param)
            if pip == 'python':
                ret=apiBase.run_python_file(url,param)
            if pip == 'code':
                func=apiBase.str_func(url)
                ret=func(param)
                try:
                    ret=ret.string
                except Exception as e:
                    pass
            if pip == 'print':
                ret=param
                print(f"param={param}")
            return ret
        if desc is None:
            desc=name
        return StructuredTool.from_function(
            company_tool,
            name=name,
            description=f"Information about {desc}",
        )
    # 注册agent
    def agent_reg(self,name,agent):
        if not name:
            return "agent name is None"
        if not agent:
            return "agent  is None"
        self.agent_list[name]=agent
        return "Agent registration successful"
    # 注册agent
    def agent_find(self,name):
        if not name:
            return "agent name is None"
        if name in self.agent_list:
            return self.agent_list[name]        
        return None
    
    # 把dir目录下的python全部转成成数组
    def dir_func(self,dataDir):
        self.funs_run={} 
        folder_names = os.listdir(dataDir)
        if dataDir.endswith('/'):
            dataDir = dataDir[:-1]
        #names=['grp_','nlp_','pic_','rag_','act_','tol','ept','ept','rft','sft','dat','ml_']    
        with open('./info.py', "r") as file:
            content = file.read()             
            fn=self.str_func(content)            
            self.funs_run['info']=fn
                            
        names=['llm','rag_','tol_','act_','rft_','grp_','ml_']
        for name in folder_names:            
            fname=dataDir + "/" + name
            if not os.path.isfile(fname):
                continue
            ls=name.split(".")
            if len(ls) !=2:
                continue
            file_extension = ls[1]            
            if file_extension != "py":
                continue
            print(f'load {name}')
            for nm in names:                
                if name.startswith(nm):                    
                    with open(fname, "r") as file:
                        content = file.read()             
                        fn=self.str_func(content)            
                        if fn:
                            self.funs_run[ls[0]]=fn
                    break
        return self.funs_run
            
    # 字符串转成class
    def argv_cls(self,funs_param,index,defVal=None):
        try:
            if len(funs_param) <= index:
                cls=defVal
            else:
                cls=self.argv(funs_param,index,defVal)        
            # 提取函数名称
            pattern = r'class\s+(\w+)\('
            match = re.search(pattern, cls) 
            name = match.group(1)                
            exec(cls)
            return eval(f'{name}')
        except Exception as e:
            return None
                
    def argv_array(self,funs_param,index,defVal=None,sep=';'):
        try:
            if len(funs_param) <= index:
                data=defVal
            else:
                data=self.argv(funs_param,index,defVal)    
            return data.rstrip(sep).split(sep)
        except Exception as e:
            return []        
    
    def cut_all(self,maxtoken,code,lsExpress,reg1,reg2,blockname):
        sql_statements = re.findall(reg1, code, re.DOTALL)
        # 防止无限的递归
        if len(sql_statements) == 1:
            if reg2 is  None:
                return code
            sql_statements = re.findall(reg2, code, re.DOTALL)
            if len(sql_statements) == 1:
                return code
        for stmt in sql_statements:
            if len(stmt) > maxtoken:
                stmt=self.cut_all(maxtoken,stmt,lsExpress,reg1,reg2,blockname)
            lsExpress.append(stmt)
        for i in range(len(lsExpress)):
            code=code.replace(lsExpress[i],f"{blockname}{i:05}")
        return code

    def save_log(self,stage,cnt):
        con=self.pool.getconn()  
        try:
            with con.cursor() as cur:
                sql = f"insert into aigc_log(stage,cnt)VALUES('{stage}','{cnt}')"
                cur.execute(sql)
                # 提交事务
                con.commit()
        finally:
            self.pool.putconn(con)
    def log(self,val):
        logging.info(val)
        
    def prop_read(self,key):
        return os.environ[key]
    
    def prop_write(self,key,val):
        os.environ[key] = val
        return True
    def encode64(self,cnt):
        if cnt is None:
            return ""
        if len(cnt) == 0:
            return ""
        encoded_str = base64.b64encode(cnt.encode('utf-8')).decode()
        return str(encoded_str)
    
    def decode64(self,cnt):
        if cnt is None:
            return ""
        if len(cnt) == 0:
            return ""
        try:
            decoded_str = base64.b64decode(cnt).decode()
        except Exception as e:
            return cnt
        return decoded_str            
    
    # 用于python组件之间的调用:apiBase.run_python_file('./set_env.py',"1" , "2","3")
    def run_python_file(self,temp_file_name,*args,check=False):
        try:
            # 检查是否已经运行    
            if check:
                # 使用 subprocess.run 执行命令            
                grep_command = 'ps -ef | grep {}'.format(temp_file_name)
                # 使用subprocess.run执行命令，并捕获输出
                result = subprocess.run(grep_command, shell=True, stdout=subprocess.PIPE, text=True)
                # 将输出分行并返回
                lines = result.stdout.splitlines()
                # 
                if len(lines) > 2:
                    return 0
                #cmd=f"nohup python3 {temp_file_name}.py > {temp_file_name}.log 2>&1 &"
                #result = os.system(cmd)
                subprocess.Popen(["nohup","python3",temp_file_name+".py",">" ,temp_file_name+".log","2>&1","&"], start_new_session=True)
                return result
            else:
                result = subprocess.run(["python3",temp_file_name+".py"  ,*args], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            # 如果脚本成功执行，返回码为0
            if result.returncode == 0:
                return result.stdout.decode()
        except Exception as e:
            print(f"[error]:{e}")
        return 0
            
    def read_json(self,dir,api_name):
        filename = api_name+'.json'        
        data = None
        file_path = os.path.join(dir,filename)
        if not os.path.exists(file_path):
            return "json not found"
        # 读取json文件        
        with open(file_path, 'r', encoding='utf-8') as file:
            data = json.load(file)
        # 添加json文件内容
        data['prompt'] = self.read_md(dir,"prompt") 
        return data
    
    def generate_md5(self,message):
        if message is None:
            return ""        
        return hashlib.md5(message.encode('utf-8')).hexdigest()
        
    
    def read_md(self,dir,api_name):
        filename = api_name+'.md'
        text = None
        file_path = os.path.join(dir,filename)
        if not os.path.exists(file_path):
            return "md not found"
        # 读取md文件        
        with open(file_path, 'r', encoding='utf-8') as file:
            text = file.read()
        return text
    
    # 调用prompt的odata接口，根据key和version从数据库读取prompt
    def odata_read(self,base_url):
        #base_url= f"{self.api_url}/crud.svc/LlmPromptTabs?$filter=pkey eq '{key}'&$select=content"
        #base_url = "http://8.130.66.117:18518/crud.svc/LlmPromptTabs?$filter=pkey%20eq%20%27test%27&$select=content"
        headers = {'Content-Type': 'application/json'}
        response = requests.get(base_url,headers=headers)
        ret = response.json()
        values = ret.get('value')
        if values is None:
            return None
        length = len(values)
        if length == 0:
            return None
        return values
    
    # 调用prompt的odata接口，根据key和version修改prompt，
    def odata_add(self,tabname,data,key):
        base_url= f"{self.api_url}crud.svc/{tabname}"
        query_url= f"{base_url}?$filter={key}"
        old_cnt = self.odata_read(query_url)
        headers = {'Content-Type': 'application/json'}
        if old_cnt is  None:            
            response = requests.post(base_url,json=data,headers=headers)
        else:
            response = requests.put(query_url,json=data,headers=headers)    
        return response.ok
    
    def prompt_read(self,pkey):
        con=self.pool.getconn()  
        try:
            with con.cursor() as cur:
                sql = f"select content from  aigc_prompt where pkey='{pkey}'"
                cur.execute(sql)
                results = cur.fetchall()
                for row in results:
                    content, = row           
                    cnt = self.decode64(content)       
                    return cnt
        finally:
            self.pool.putconn(con)
            
    def prompt_add(self,pkey,prompt):
        content = self.encode64(prompt)
        con=self.pool.getconn()  
        try:
            with con.cursor() as cur:
                sql = f"select pkey from  aigc_prompt where pkey='{pkey}'"
                cur.execute(sql)
                results = cur.fetchall()            
                if results:
                    sql = f"update  aigc_prompt set content='{content}' where pkey='{pkey}'"
                else:
                    sql = f"insert into aigc_prompt(pkey,content)values('{pkey}','{content}')"
                cur.execute(sql)
                con.commit()
        finally:
            self.pool.putconn(con)
        return 'ok'
    
    def summary_read(self,abstract):
        md5 = self.generate_md5(abstract)
        base_url= f"{self.api_url}crud.svc/LlmAbstractTabs?$filter=summarize eq '{md5}'&$select=path"
        values=self.odata_read(base_url)
        if values is None:
            return None
        path = values[0]
        # 读取文件内容        
        with open(path, 'r', encoding='utf-8') as file:
            text = file.read()
        return text
    
    def call_api(self,api_name,params,judge=False):   
        if api_name is None:
            return None
        print(f"[call_api]api={api_name};params={params}") 
        #去掉多余的数据
        if params:
            params = params.split("\n")[0]
        headers = {
        'Accept': 'application/json',
        'Content-Type': 'application/json'
        }        
        response = requests.post(api_name, {"cnt":params}, headers=headers)
        response.encoding = 'utf-8'
        ret = response.json()
        if judge:
            if 'data' not in ret:
                return None                   
        return ret
    
    def call_pip(self,api_name,params): 
        if api_name is None:
            return None
        if not api_name.endswith(".hpl"):
            api_name=api_name+".hpl"
        headers = {
        'Accept': 'application/json; charset=utf-8',        
        'Content-Type': 'application/json; charset=utf-8'
        }
        if params is None:
            return None
        if api_name.find('$PROJECT_HOME') != -1:
            template = Template(api_name)
            api_name = template.substitute(PROJECT_HOME=os.environ["PROJECT_HOME"])
        #去掉多余的数据
        params = params.split("\n")[0]
        pos = api_name.find("/samples/")
        if pos != -1:
            pos += len("/samples/")
            api_name = api_name[pos:]        
        #http://8.130.66.117:8082/88api/webssh/echo.hpl?cnt=%7B%22cnt%22:%225L2g5aW9%22%7D
        base_url= f'{os.environ["LLMDB_SVR"]}/{api_name}?cnt=%7B%22cnt%22:%22{params}%22%7D'
        print("[call_api]base_url=",base_url)
        response = requests.post(base_url, headers=headers)
        response.encoding = 'utf-8'
        ret= response.json()
        if 'data' not in ret:
            return None          
        return ret

    # 提供系统命令行，如打开drawio，打开svg，播放
    def call_cmd(self,cls):
        retInf =""
        # 定义shell脚本的路径
        script_path = cls.get('api-param')
        cmd = cls.get('param')
        basename, extension = os.path.splitext(script_path)
        if extension =='.py':
            script_path="python "+ script_path
            result = subprocess.run(script_path, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
            retInf=f"{result.stdout} {result.stderr}"
        
        if extension =='.sh':
            script_path="/bin/bash "+ script_path
            result = subprocess.run(script_path, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
            retInf=f"{result.stdout} {result.stderr}"
        
        script_path=cmd +" " + script_path
        print("script_path="+script_path)
        result = subprocess.run(script_path, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        retInf=f"{result.stdout} {result.stderr}"
        return retInf
    
    # 向量化
    def vector_embed(self,docs:list):
        return emObj.call(docs)            
    
    # 向量化
    def vector_embed_api(self,docs:list):
        try:
            headers = {
                #'Authorization': f'Bearer {api_key}',
                'Content-Type': 'application/json'
            }            
            response = requests.post("http://127.0.0.1:51818/embeddings", headers=headers,json=docs)
            resp=response.json()           
            return resp
        except Exception as e:
            print(e)            
        return None
    
    
    def llm_chat(self,sys_prompt,usr_prompt,code=None,error=None):
        if sys_prompt is None:
            sys_prompt="You are a helpful assistant"
        if usr_prompt is None:
            usr_prompt="hi"
        if code is None:
            code=""
        if error is None:
            error=""     
        prompt=""
        if code.strip():
            prompt="##code\r\n"+code+"\r\n"
        if error.strip():
            prompt = prompt + "##error\r\n The error is as follows:"+error+"\r\n" 
        prompt = prompt + "##human\r\n"+usr_prompt            
        
        api=os.environ["OPENAI_API_BASE"]+"chat/completions"
        api_key=os.environ["OPENAI_API_KEY"]
        model_name=os.environ["OPENAI_MODEL"]
        max_tokens=os.environ["OPENAI_API_MAX_TOKENS"]
        temperature=os.environ["OPENAI_API_TEMPERATURE"]
        headers = {
            'Authorization': f'Bearer {api_key}',
            'Content-Type': 'application/json'
        }
        api = api.strip()
        token=int(eval(max_tokens))
        tmt=float(temperature)
        data = {
            'model': model_name, 
            'messages': [
                {"role": "system", "content": sys_prompt},
                {"role": "user", "content":prompt}
            ],
            "temperature" : tmt, 
            "max_tokens" : token ,
            "stream": False
        }
        
        # 发送POST请求到OpenAI的Chat模型接口
        #self.log("api="+api)        
        try:
            response = requests.post(api, headers=headers, json=data)
            resp=response.json()            
            if 'choices' in resp:
                resp=resp['choices'][0]
            ret=resp['message']['content']
        except Exception as e:
            #print(e)
            ret=f"error={resp['error']}" 
        return ret    

apiBase = ApiBase()

#apiBase.run_python_file("ApiVctSvr",check=True)
#apiBase.vector_embed(["hi"])
#apiBase.get_random(10,10)
# print(apiBase.call_api('http://8.130.66.117:8082/88api/langgraph/echoA.hpl',''))
# print(apiBase.call_pip("$PROJECT_HOME/88api/webssh/echo","5L2g5aW9"))
# print(apiBase.call_pip("/88api/langgraph/echoA",""))

# str= apiBase.encode64("你好")
# print(str)
# code = apiBase.decode64(str)
# if code == "你好":
#     print("ok")
#apiBase.run_python("print('hello world')")
#print(apiBase.prompt_read("test"))
#apiBase.summary_add("abcd","E:/04robot/llm-ssh/llm-ssh/webssh/PipAgent/tools/createdoc.hpl")
#apiBase.summary_read("abcd")